import json
import os
import sys

sys.path.append(os.path.dirname(sys.path[0]))
from aws_model.athena_submit import submit
from project_utils.import_excel_notime import import_excel
from aws_model.aws_stepfunc import send_task_failure, send_task_success
from base_class.inspect_logging import print_to_log
from aws_model.aws_s3 import Operation_aws_s3

'''
传入参数:
{
 "source":源文件地址，定位到具体文件, C://Users/DELL/Desktop/虹美/快益点压缩文件
 "target":目标地址，定位到文件夹，需要带文件类型，比如本地文件系统需要前缀 file:// s3需要 s3://,
 "sheets":excel需要导入的sheet名，多个sheet 逗号隔开,
 "type":映射类型，对应import_mapping表中的 type，现有 kyd,ac365,
 "version":版本，对应import_mapping表中version，用于过滤,
}
'''
if __name__ == '__main__':
    input = json.loads(sys.argv[1])
    tasktoken = input.get('TaskToken')
    if 'env' in input:
        env = input['env']
    else:
        env = 'prod'

    if env == 'prod':
        root_path = 's3://basic-data/ac_kyd/increment/'
        table_name = 'original.ac_kyd_all'
        target_root_path = "s3://original/ac_kyd_all"
    else:
        root_path = 's3://basic-data/ac_kyd/increment-test/'
        table_name = 'test_group.original_ac_kyd_all'
        target_root_path = "s3://testgroup-data/original/ac_kyd_all"

    s3 = Operation_aws_s3()
    sources = []
    if 'source' in input and input['source']:
        sources = str(input['source']).split(',')
    else:
        dirs = s3.ls(root_path)
        for obj in dirs:
            if str(obj['object']).endswith('.xls') or str(obj['object']).endswith('.xlsx'):
                sources.append(f"s3://{obj['bucket']}/{obj['object']}")

    print_to_log(sources)
    try:
        for source in sources:
            if 'version' not in input or not input['version']:
                version = 'default'
            else:
                version = input['version']

            i = source.rfind('/')
            partition = source[i:][1:8]
            target = f'{target_root_path}/import_date={partition}/'
            sheets = '维修结算'

            print(source, target, sheets, input['type'], version)
            import_excel(source, target, sheets, 'kyd', version)
            submit(f"alter table {table_name} add if not exists partition (import_date = '{partition}')")
            s3.rename(source,f'{source}.complete')
            print_to_log(f'文件{source}导入完成')
        output = {"msg": f"文件导入完成", "Payload": {}}
        output = json.dumps(output)
        step_response = send_task_success(tasktoken, output)
    except Exception as e:
        print_to_log("任务失败")
        print_to_log("原因： ", str(e))
        send_task_failure(str(e), str(e), tasktoken)
